#include <iostream>
#include <string>
#include <ctime>
#include <fstream>
#include <vector>
#include <thread>
#include <mutex>
#include <winsock2.h>
#include <pcap.h>

#pragma comment(lib, "wpcap.lib")
#pragma comment(lib, "Ws2_32.lib")

#define LOG_FILE "D:\\WpdPack\\project\\Project1\\x64\\Debug\\ftp_logins.csv"

// FTP登录记录结构体
struct FTPLogin {
    std::string timestamp;
    std::string src_mac;
    std::string src_ip;
    std::string dst_mac;
    std::string dst_ip;
    std::string username;
    std::string password;
    std::string status;
};

// 全局变量（线程安全控制）
std::mutex log_mutex;
FTPLogin current_login;

// 写入CSV日志
void write_log(const FTPLogin& login) {
    std::lock_guard<std::mutex> lock(log_mutex);
    std::ofstream file(LOG_FILE, std::ios::app);
    if (file.is_open()) {
        file << login.timestamp << ","
            << login.src_mac << ","
            << login.src_ip << ","
            << login.dst_mac << ","
            << login.dst_ip << ","
            << login.username << ","
            << login.password << ","
            << login.status << "\n";
        file.close();
    }
}

// 解析以太网帧
void parse_eth_header(const u_char* pkt_data, std::string& src_mac, std::string& dst_mac) {
    char buffer[18];
    const u_char* src = pkt_data + 6;
    const u_char* dst = pkt_data;
    sprintf_s(buffer, sizeof(buffer), "%02X:%02X:%02X:%02X:%02X:%02X",
        dst[0], dst[1], dst[2], dst[3], dst[4], dst[5]);
    dst_mac = buffer;
    sprintf_s(buffer, sizeof(buffer), "%02X:%02X:%02X:%02X:%02X:%02X",
        src[0], src[1], src[2], src[3], src[4], src[5]);
    src_mac = buffer;
}

// 解析IP头部
void parse_ip_header(const u_char* pkt_data, std::string& src_ip, std::string& dst_ip) {
    const struct ip_header {
        u_char  ihl : 4;
        u_char  version : 4;
        u_char  tos;
        u_short tot_len;
        u_short id;
        u_short frag_off;
        u_char  ttl;
        u_char  protocol;
        u_short check;
        u_int   saddr;
        u_int   daddr;
    }*ip = reinterpret_cast<const ip_header*>(pkt_data + 14); 

    char src[16], dst[16];
    inet_ntop(AF_INET, &ip->saddr, src, sizeof(src));
    inet_ntop(AF_INET, &ip->daddr, dst, sizeof(dst));
    src_ip = src;
    dst_ip = dst;
}

// 解析FTP载荷
void parse_ftp_payload(const u_char* payload, int len, FTPLogin& login) {
    std::string data(reinterpret_cast<const char*>(payload), len);
    size_t pos;

    if ((pos = data.find("USER ")) != std::string::npos) {
        login.username = data.substr(pos + 5, data.find("\r\n", pos) - pos - 5);
    }
    else if ((pos = data.find("PASS ")) != std::string::npos) {
        login.password = data.substr(pos + 5, data.find("\r\n", pos) - pos - 5);
    }
    else if (data.find("230 ") != std::string::npos) {
        login.status = "SUCCEED";
        write_log(login);
        current_login = FTPLogin(); 
    }
    else if (data.find("530 ") != std::string::npos) {
        login.status = "FAILED";
        write_log(login);
        current_login = FTPLogin();
    }
}

// 抓包回调函数
void packet_handler(u_char* param, const pcap_pkthdr* header, const u_char* pkt_data) {
    // 解析时间戳
    time_t raw_time = header->ts.tv_sec;
    struct tm timeinfo;
    localtime_s(&timeinfo, &raw_time);
    char timestamp[20];
    strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", &timeinfo);

    // 解析MAC和IP
    std::string src_mac, dst_mac, src_ip, dst_ip;
    parse_eth_header(pkt_data, src_mac, dst_mac);
    parse_ip_header(pkt_data, src_ip, dst_ip);

    // 解析TCP端口
    const struct tcp_header {
        u_short sport;
        u_short dport;
        u_int   seq;
        u_int   ack;
        u_char  data_offset;
        u_char  flags;
        u_short window;
        u_short check;
        u_short urg_ptr;
    }*tcp = reinterpret_cast<const tcp_header*>(pkt_data + 14 + 20);

    if (ntohs(tcp->dport) == 21) {
        int tcp_header_len = (tcp->data_offset >> 4) * 4;
        const u_char* payload = pkt_data + 14 + 20 + tcp_header_len;
        int payload_len = header->len - (14 + 20 + tcp_header_len);

        // 填充当前登录信息
        current_login.timestamp = timestamp;
        current_login.src_mac = src_mac;
        current_login.src_ip = src_ip;
        current_login.dst_mac = dst_mac;
        current_login.dst_ip = dst_ip;

        parse_ftp_payload(payload, payload_len, current_login);
    }
}

int main() {
    pcap_if_t* alldevs;
    char errbuf[PCAP_ERRBUF_SIZE];

    // 获取网卡列表
    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        std::cerr << "Error finding devices: " << errbuf << std::endl;
        return 1;
    }

    // 选择第一个网卡
    pcap_t* handle = pcap_open_live(alldevs->name, 65535, 1, 1000, errbuf);
    if (!handle) {
        std::cerr << "Error opening device: " << errbuf << std::endl;
        return 1;
    }

    // 设置过滤器（仅捕获目标端口21的TCP流量）
    struct bpf_program fp;
    std::string filter = "tcp dst port 21";
    pcap_compile(handle, &fp, filter.c_str(), 0, 0);
    pcap_setfilter(handle, &fp);

    // 开始抓包
    pcap_loop(handle, 0, packet_handler, nullptr);

    pcap_close(handle);
    pcap_freealldevs(alldevs);
    return 0;
}